package cn.git.agent.loader;

import cn.git.agent.PluginBootstrap;
import cn.git.agent.boot.AgentPackagePath;
import lombok.RequiredArgsConstructor;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.ArrayUtils;

import java.io.*;
import java.net.URL;
import java.util.*;
import java.util.concurrent.locks.ReentrantLock;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

/**
 * @description: 用于加载插件和插件的拦截器
 * @program: bank-credit-sy
 * @author: lixuchun
 * @create: 2024-12-26
 */
public class AgentClassLoader extends ClassLoader {

    /**
     * 默认的类加载器
     * 用于定义 Mysql5Instrumentation等类，除了插件的interceptor
     */
    private static AgentClassLoader DEFAULT_CLASS_LOADER;

    /**
     * 插件的类路径
     */
    private List<File> classPath;
    private List<Jar> allJars;

    /**
     * 加载jar文件的锁，避免多次加载
     */
    private ReentrantLock jarScanLock = new ReentrantLock();

    /**
     * 构造函数
     *
     * @param parent
     */
    public AgentClassLoader(ClassLoader parent) {
        super(parent);
        // agent.jar的目录
        File agentJarDir = AgentPackagePath.getPath();
        classPath = new LinkedList<>();
        classPath.add(new File(agentJarDir, "plugins"));
    }

    public static AgentClassLoader getDefaultClassLoader() {
        return DEFAULT_CLASS_LOADER;
    }

    /**
     * 初始化默认的类加载器
     */
    public static void initDefaultClassLoader() {
        if (DEFAULT_CLASS_LOADER == null) {
            DEFAULT_CLASS_LOADER = new AgentClassLoader(PluginBootstrap.class.getClassLoader());
        }
    }

    /**
     * 重写findClass方法
     *  loadClass -> 自己回调 findClass（自己定义自己的类加载器） -> defineClass
     * @param name
     * @return
     * @throws ClassNotFoundException
     */
    @Override
    protected Class<?> findClass(String name) throws ClassNotFoundException {
        List<Jar> allJars = getAllJars();

        // 获取类名
        String path = name.replace(".", "/").concat(".class");

        for (Jar jar : allJars) {
            JarEntry entry = jar.jarFile.getJarEntry(path);
            if (entry == null) {
                continue;
            }

            // 获取jar文件
            try {
                String jarPath = jar.sourceFile.getAbsolutePath().replace('\\', '/');
                URL url = new URL("jar:file:/" + jarPath + "!/" + path);
                byte[] bytes = readBytesFromURL(url);
                return defineClass(name, bytes, 0, bytes.length);
            } catch (Exception e) {
                e.printStackTrace();
                System.out.println("Can not load jar file:" + jar.sourceFile.getAbsolutePath());
            }
        }
        throw new ClassNotFoundException("can not find " + name);
    }

    private byte[] readBytesFromURL(URL url) throws IOException {
        try (InputStream is = url.openStream();
             ByteArrayOutputStream buffer = new ByteArrayOutputStream()) {
            int nRead;
            byte[] data = new byte[16384];
            while ((nRead = is.read(data, 0, data.length)) != -1) {
                buffer.write(data, 0, nRead);
            }
            buffer.flush();
            return buffer.toByteArray();
        }
    }

    /**
     * 重写getResources方法
     *
     * @param name
     * @return
     * @throws IOException
     */
    @Override
    public Enumeration<URL> getResources(String name) throws IOException {
        List<URL> allResources = new LinkedList<>();
        List<Jar> allJars = getAllJars();
        for (Jar jar : allJars) {
            JarEntry entry = jar.jarFile.getJarEntry(name);
            if (entry != null) {
                String jarPath = jar.sourceFile.getAbsolutePath().replace('\\', '/');
//                allResources.add(new URL("jar:file:" + jar.sourceFile.getAbsolutePath() + "!/" + name));
                allResources.add(new URL("jar:file:" + jarPath + "!/" + name));
            }
        }

        final Iterator<URL> iterator = allResources.iterator();
        return new Enumeration<URL>() {
            @Override
            public boolean hasMoreElements() {
                return iterator.hasNext();
            }

            @Override
            public URL nextElement() {
                return iterator.next();
            }
        };
    }

    /**
     * 重写getResource方法
     *
     * @param name
     * @return
     */
    @Override
    public URL getResource(String name) {
        for (Jar jar : allJars) {
            JarEntry jarEntry = jar.jarFile.getJarEntry(name);
            if (jarEntry != null) {
                try {
                    return new URL("jar:file" + jar.sourceFile.getAbsolutePath() + "!/" + name);
                } catch (Exception e) {
                    System.out.println("Can not load jar file:" + jar.sourceFile.getAbsolutePath());
                    e.printStackTrace();
                }
            }
        }
        return null;
    }

    /**
     * 获取所有的jar文件
     * @return
     */
    private List<Jar> getAllJars() {
        if (allJars == null) {
            jarScanLock.lock();
            try {
                // 避免多次加载
                if (allJars == null) {
                    allJars = doGetJar();
                }
            } finally {
                jarScanLock.unlock();
            }
        }
        return allJars;
    }

    /**
     * 获取所有的jar文件
     *
     * @return
     */
    private List<Jar> doGetJar() {
        // 返回参数
        List<Jar> list = new LinkedList<>();
        //eg: D:/apm-sniffer/apm-agent-core/target/classes/plugins
        for (File path : classPath) {
            if (path.exists() && path.isDirectory()) {
                // jarFile名称
                String[] jarFileNames = path.list(new FilenameFilter() {
                    @Override
                    public boolean accept(File dir, String name) {
                        return name.endsWith(".jar");
                    }
                });
                if (ArrayUtils.isEmpty(jarFileNames)) {
                    continue;
                }

                for (String jarFileName : jarFileNames) {
                    File jarSourceFile = new File(path, jarFileName);
                    Jar jar = null;
                    try {
                        jar = new Jar(new JarFile(jarSourceFile), jarSourceFile);
                        list.add(jar);
                        System.out.println("Load jar file:" + jarSourceFile.getAbsolutePath());
                    } catch (IOException e) {
                        e.printStackTrace();
                        System.out.println("Can not load jar file:" + jarSourceFile.getAbsolutePath());
                    }
                }
            }
        }
        return list;
    }

    /** 
     * @description: 插件的类加载器
     * @program: bank-credit-sy
     * @author: lixuchun
     * @create: 2024-12-26
     */
    @RequiredArgsConstructor
    private static class Jar {
        /**
         * jar文件
         */
        private final JarFile jarFile;

        /**
         * jar文件对应的目录
         */
        private final File sourceFile;

    }

}
